import { _decorator, CCString, Enum, Vec3, Node, Layers, Color, CCBoolean, CCFloat, CCInteger, CCClass } from 'cc';
import { SensorComponent } from './SensorComponent';
import { StackingSensor } from './StackingSensor';
import { SensorCompressionType } from './CompressionSpec';
import { ISensor } from './ISensor';
import { GridSensorBase } from './GridSensorBase';
import { IGridPerception } from './IGridPerception';
import { BoxOverlapChecker } from './BoxOverlapChecker';
import { CCCAgentsError } from '../CCCAgentsError';
import { OneHotGridSensor } from './OneHotGridSensor';
const { ccclass, property, menu } = _decorator;

Enum(SensorCompressionType);

@ccclass('GridSensorComponent')
@menu('ML Agents/Grid Sensor')
export class GridSensorComponent extends SensorComponent {
	_debugSensor: GridSensorBase = null!;

	_sensors: GridSensorBase[] = [];

	_gridPerception: IGridPerception = null!;

	@property(CCString)
	sensorName: string = 'GridSensor';

	@property({ serializable: true, visible: true })
	private _cellScale: Vec3 = new Vec3(1, 0.01, 1);
	public get cellScale(): Vec3 {
		return this._cellScale;
	}
	public set cellScale(value: Vec3) {
		this._cellScale = value;
	}

	@property({ serializable: true, visible: true })
	private _gridSize: Vec3 = new Vec3(16, 1, 16);
	public get gridSize(): Vec3 {
		return this._gridSize;
	}
	public set gridSize(value: Vec3) {
		if (value.y != 1) {
			this._gridSize = new Vec3(value.x, 1, value.z);
		} else {
			this._gridSize = value;
		}
	}

	@property(CCBoolean)
	rotateWithAgent: boolean = false;

	@property(Node)
	private _agentGameObject: Node = null!;
	public get agentGameObject(): Node {
		return this._agentGameObject === null ? this.node : this._agentGameObject;
	}
	public set agentGameObject(value: Node) {
		this._agentGameObject = value;
	}

	@property([CCString])
	detectableTags: string[] = [];

	@property({ type: Layers.BitMask })
	colliderMask: Layers.BitMask = Layers.BitMask.NONE;

	@property({ range: [1, 50, 1] })
	public observationStacks = 1;

	@property({ type: SensorCompressionType, serializable: true, visible: true })
	private _compression: SensorCompressionType = SensorCompressionType.PNG;
	public get compression(): SensorCompressionType {
		return this._compression;
	}
	public set compression(value: SensorCompressionType) {
		this._compression = value;
		this.updateSensor();
	}

	@property(CCInteger)
	initialColliderBufferSize: number = 4;

	@property(CCInteger)
	maxColliderBufferSize: number = 500;

	@property(CCBoolean)
	showGizmos: boolean = false;

	@property([Color])
	debugColors: Color[] = [];

	@property(CCInteger)
	gizmoYOffset: number = 0;

	createSensors(): ISensor[] {
		this._gridPerception = new BoxOverlapChecker(
			this.cellScale,
			this.gridSize,
			this.rotateWithAgent,
			this.colliderMask,
			this.node,
			this.agentGameObject,
			this.detectableTags,
			this.initialColliderBufferSize,
			this.maxColliderBufferSize
		);

		this._debugSensor = new GridSensorBase(
			'debugGridSensor',
			this.cellScale,
			this.gridSize,
			this.detectableTags,
			SensorCompressionType.None
		);

		this._sensors = this.getGridSensors();

		if (this._sensors.length < 1) {
			throw new CCCAgentsError('No Grid Sensors found');
		}

		this._sensors[0]._gridPerception = this._gridPerception;
		this._sensors.forEach((sensor) => {
			this._gridPerception.registerSensor(sensor);
		});

		if (this.observationStacks != 1) {
			const sensors = new Array(this._sensors.length);
			for (let i = 0; i < sensors.length; i++) {
				sensors[i] = new StackingSensor(this._sensors[i], this.observationStacks);
			}
			return sensors;
		} else {
			return this._sensors;
		}
	}

	private updateSensor() {
		if (this._sensors != null) {
			this._gridPerception.rotateWithAgent = this.rotateWithAgent;
			this._gridPerception.colliderMask = this.colliderMask;

			this._sensors.forEach((sensor) => {
				sensor.compressionType = this._compression;
			});
		}
	}

	protected onDestroy(): void { }

	/**
	 * @en
	 * Get an array of GridSensors to be added in this component.
	 * Override this method and return custom GridSensor implementations.
	 *
	 * @zh
	 * 获取要添加到此组件中的 GridSensors 数组。
	 * 重写此方法并返回自定义 GridSensor 实现。
	 *
	 *
	 * @returns
	 */
	getGridSensors(): GridSensorBase[] {
		const sensorList = new Array<GridSensorBase>();
		const sensor = new OneHotGridSensor(
			this.sensorName + '_OneHot',
			this.cellScale,
			this.gridSize,
			this.detectableTags,
			this._compression
		);
		sensorList.push(sensor);
		return sensorList;
	}
}
